import numpy as np
from enum import Enum
from typing import Dict

class SubChannelType(Enum):
    """子通道类型枚举"""
    CORNER = "corner"    # 角通道
    EDGE = "edge"       # 边通道
    CENTER = "center"   # 中心通道

class SubChannel:
    """子通道类，定义单个子通道的几何和物理特性"""
    def __init__(self, id: int, type: SubChannelType, geometry: Dict):
        self.id = id
        self.type = type
        
        # 几何参数
        self.area = geometry['area']                          # 流通面积
        self.wetted_perimeter = geometry['wetted_perimeter']  # 湿周长
        self.heated_perimeter = geometry['heated_perimeter']  # 加热周长
        self.hydraulic_diameter = geometry['hydraulic_diameter']  # 水力直径
        self.gap_width = geometry['gap_width']                # 间隙宽度
        self.gap_length = geometry['gap_length']              # 间隙长度
        
        # 相邻子通道连接信息
        self.connections = []  # 存储相邻子通道的连接信息
        
        # 场量初始化
        self.alpha = None              # 空泡份额
        self.pressure = None           # 压力场
        self.velocity = {              # 速度场
            'liquid': {
                'axial': None,
                'lateral': None
            },
            'vapor': {
                'axial': None,
                'lateral': None
            }
        }
        
        # 保存上一时间步的值
        self.previous_alpha = None
        self.previous_pressure = None
        self.previous_velocity = {     # 上一时间步速度场
            'liquid': {
                'axial': None,
                'lateral': None
            },
            'vapor': {
                'axial': None,
                'lateral': None
            }
        }
        
        # 温度场
        self.temperature = {
            'liquid': None,
            'vapor': None
        }
        
        # 密度场
        self.density = {
            'liquid': None,
            'vapor': None
        }
        
        # TVD格式相关
        self.previous_drift_flux = None  # 上一时间步的漂移通量
        self.previous_mass_flux = {      # 上一时间步的质量通量
            'vapor': None,
            'liquid': None
        }
    
    def add_connection(self, neighbor_id: int, connection_type: str, gap_width: float):
        """添加与相邻子通道的连接信息
        
        Args:
            neighbor_id: 相邻子通道的ID
            connection_type: 连接类型 ('axial' 或 'lateral')
            gap_width: 间隙宽度
        """
        self.connections.append({
            'neighbor_id': neighbor_id,
            'type': connection_type,
            'gap_width': gap_width
        })
    
    def get_connection_parameters(self, adjacent_channel) -> Dict:
        """获取与相邻子通道的连接参数
        
        Args:
            adjacent_channel: 相邻子通道对象
            
        Returns:
            Dict: 包含连接参数的字典
        """
        # 查找连接信息
        connection = next(
            (conn for conn in self.connections 
             if conn['neighbor_id'] == adjacent_channel.id),
            None
        )
        
        if connection is None:
            raise ValueError(f"未找到与通道 {adjacent_channel.id} 的连接信息")
        
        # 计算质心距离（简化模型，使用固定值）
        if connection['type'] == 'lateral':
            centroid_distance = self.gap_width
        else:  # axial
            centroid_distance = self.gap_length
        
        return {
            'connection_type': connection['type'],
            'gap_width': connection['gap_width'],
            'gap_length': self.gap_length,
            'centroid_distance': centroid_distance
        }
    
    def get_geometry_parameters(self) -> Dict:
        """获取子通道的几何参数
        
        Returns:
            Dict: 包含所有几何参数的字典
        """
        return {
            'area': self.area,
            'wetted_perimeter': self.wetted_perimeter,
            'heated_perimeter': self.heated_perimeter,
            'hydraulic_diameter': self.hydraulic_diameter,
            'gap_width': self.gap_width,
            'gap_length': self.gap_length
        }
    
    def save_current_state(self):
        """保存当前状态用于下一时间步计算"""
        self.previous_alpha = np.copy(self.alpha)
        self.previous_pressure = np.copy(self.pressure)
        self.previous_velocity = {
            'vapor': {
                'axial': np.copy(self.velocity['vapor']['axial']),
                'lateral': np.copy(self.velocity['vapor']['lateral'])
            },
            'liquid': {
                'axial': np.copy(self.velocity['liquid']['axial']),
                'lateral': np.copy(self.velocity['liquid']['lateral'])
            }
        }
        
        # 初始化或更新TVD相关数组
        if self.previous_drift_flux is None:
            self.previous_drift_flux = np.zeros_like(self.alpha)
        if self.previous_mass_flux['vapor'] is None:
            self.previous_mass_flux['vapor'] = np.zeros_like(self.alpha)
            self.previous_mass_flux['liquid'] = np.zeros_like(self.alpha)

class Grid:
    """网格类，管理所有子通道和其连接关系"""
    def __init__(self, n_axial=20, channel_length=3.66):
        self.n_axial = n_axial  # 轴向节点数
        self.n_channels = 36     # 6x6网格的子通道总数
        self.subchannels = []    # 存储所有子通道
        self.channel_length = channel_length  # 通道长度
        self.dz = channel_length / (n_axial - 1)  # 轴向网格间距
        self.dx = 0.0126  # 横向网格间距（根据实际燃料棒间距设置）
        
        # 时间步相关
        self.current_time = 0.0  # 当前时间
        self.dt = 0.0  # 当前时间步长
        self.time_step_number = 0  # 当前时间步编号
        
        # 网格索引映射
        self.axial_indices = np.arange(n_axial)
        self.channel_indices = np.arange(self.n_channels)
        
        # 子通道类型映射
        self.channel_types = self._initialize_channel_types()
        
        # 创建所有子通道
        self._create_all_subchannels()
        
        # 建立子通道之间的连接关系
        self.setup_connections()
        
        # 初始化场变量
        self.initialize_field_variables()
        
    def _initialize_channel_types(self):
        """初始化6x6网格的子通道类型映射"""
        # 1代表角通道，2代表边通道，3代表中心通道
        channel_type_map = [
            [1, 2, 2, 2, 2, 1],
            [2, 3, 3, 3, 2, 2],
            [2, 3, 3, 3, 2, 2],
            [2, 3, 3, 3, 2, 2],
            [2, 2, 2, 2, 2, 2],
            [1, 2, 2, 2, 2, 1]
        ]
        channel_types = {}
        for i in range(6):
            for j in range(6):
                channel_id = i * 6 + j
                if channel_type_map[i][j] == 1:
                    channel_types[channel_id] = SubChannelType.CORNER
                elif channel_type_map[i][j] == 2:
                    channel_types[channel_id] = SubChannelType.EDGE
                else:
                    channel_types[channel_id] = SubChannelType.CENTER
        return channel_types
        
    def create_subchannel(self, id):
        """创建新的子通道，根据ID确定类型和几何参数"""
        channel_type = self.channel_types[id % 36]  # 使用取模运算处理多层
        
        # 根据通道类型设置几何参数
        if channel_type == SubChannelType.CORNER:
            geometry = {
                'area': 4.293e-5,              # m^2
                'wetted_perimeter': 0.01246,   # m
                'heated_perimeter': 0.00785,   # m
                'hydraulic_diameter': 0.01377,  # m
                'gap_width': 0.0033,           # m
                'gap_length': 0.0126           # m
            }
        elif channel_type == SubChannelType.EDGE:
            geometry = {
                'area': 7.447e-5,              # m^2
                'wetted_perimeter': 0.02734,   # m
                'heated_perimeter': 0.01570,   # m
                'hydraulic_diameter': 0.01089,  # m
                'gap_width': 0.0033,           # m
                'gap_length': 0.0126           # m
            }
        else:  # CENTER
            geometry = {
                'area': 1.5876e-4,             # m^2
                'wetted_perimeter': 0.02985,   # m
                'heated_perimeter': 0.02355,   # m
                'hydraulic_diameter': 0.02125,  # m
                'gap_width': 0.0033,           # m
                'gap_length': 0.0126           # m
            }
        
        subchannel = SubChannel(id, channel_type, geometry)
        self.subchannels.append(subchannel)
        return subchannel
    
    def setup_connections(self):
        """建立子通道之间的连接关系"""
        for subchannel in self.subchannels:
            base_id = subchannel.id % 36  # 当前层内的ID
            row = base_id // 6
            col = base_id % 6
            
            # 轴向连接
            if subchannel.id < self.n_channels * (self.n_axial - 1):
                subchannel.add_connection(
                    subchannel.id + 36,  # 连接到上一层对应位置
                    'axial',
                    self.dz
                )
            
            # 横向连接
            # 连接右侧子通道
            if col < 5:  # 不是最右列
                subchannel.add_connection(
                    subchannel.id + 1,
                    'lateral',
                    self.dx
                )
            
            # 连接下方子通道
            if row < 5:  # 不是最底行
                subchannel.add_connection(
                    subchannel.id + 6,
                    'lateral',
                    self.dx
                )
    
    def initialize_field_variables(self):
        """初始化场变量"""
        for subchannel in self.subchannels:
            # 初始化空泡份额
            subchannel.alpha = np.zeros(self.n_axial)
            subchannel.previous_alpha = np.zeros(self.n_axial)
            
            # 初始化压力场
            subchannel.pressure = np.ones(self.n_axial) * 1.0e5  # 假设初始压力为1bar
            subchannel.previous_pressure = np.ones(self.n_axial) * 1.0e5
            
            # 初始化速度场
            for phase in ['liquid', 'vapor']:
                subchannel.velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.temperature[phase] = np.ones(self.n_axial) * 300.0  # 假设初始温度为300K
                subchannel.density[phase] = np.zeros(self.n_axial)
    
    def set_boundary_conditions(self, inlet_conditions, outlet_conditions):
        """设置边界条件"""
        # 入口边界条件
        for subchannel in self.subchannels[:36]:  # 第一层的所有子通道
            # 设置速度
            for phase in ['liquid', 'vapor']:
                subchannel.velocity[phase]['axial'][0] = inlet_conditions[phase]['velocity']
                # 使用全局温度
                subchannel.temperature[phase][0] = inlet_conditions['temperature']
            
            # 设置压力和空泡份额
            subchannel.pressure[0] = inlet_conditions['pressure']
            subchannel.alpha[0] = inlet_conditions['alpha']
        
        # 出口边界条件
        for subchannel in self.subchannels[-36:]:  # 最后一层的所有子通道
            subchannel.pressure[-1] = outlet_conditions['pressure']
            
    def advance_time(self, dt):
        """推进时间步"""
        self.current_time += dt
        self.time_step_number += 1
        self.dt = dt
        
        # 保存当前状态为上一时间步的值
        for subchannel in self.subchannels:
            subchannel.save_current_state() 
    
    def _create_all_subchannels(self):
        """创建所有子通道实例"""
        # 清空现有子通道列表
        self.subchannels = []
        
        # 为每个轴向层创建子通道
        for k in range(self.n_axial):
            for i in range(36):  # 6x6网格
                channel_id = k * 36 + i  # 计算全局ID
                self.create_subchannel(channel_id)
    
    def get_adjacent_channels(self, channel_id):
        """获取指定通道的相邻通道
        Args:
            channel_id: 子通道ID
        Returns:
            adjacent_channels: 相邻通道列表
        """
        # 根据通道类型确定相邻通道
        channel = self.subchannels[channel_id]
        adjacent_ids = []
        base_id = channel_id % 36  # 获取当前层内的ID
        
        # 根据通道位置确定相邻通道
        if channel.type == SubChannelType.CORNER:
            # 角通道有2个相邻通道
            if base_id == 0:  # 左上角
                adjacent_ids = [base_id + 1, base_id + 6]
            elif base_id == 5:  # 右上角
                adjacent_ids = [base_id - 1, base_id + 6]
            elif base_id == 30:  # 左下角
                adjacent_ids = [base_id + 1, base_id - 6]
            elif base_id == 35:  # 右下角
                adjacent_ids = [base_id - 1, base_id - 6]
        
        elif channel.type == SubChannelType.EDGE:
            # 边通道有3个相邻通道
            row = base_id // 6
            col = base_id % 6
            
            if row == 0:  # 上边
                adjacent_ids = [base_id - 1, base_id + 1, base_id + 6]
            elif row == 5:  # 下边
                adjacent_ids = [base_id - 1, base_id + 1, base_id - 6]
            elif col == 0:  # 左边
                adjacent_ids = [base_id + 1, base_id - 6, base_id + 6]
            elif col == 5:  # 右边
                adjacent_ids = [base_id - 1, base_id - 6, base_id + 6]
        
        else:  # CENTER
            # 中心通道有4个相邻通道
            row = base_id // 6
            col = base_id % 6
            adjacent_ids = [
                base_id - 1,  # 左
                base_id + 1,  # 右
                base_id - 6,  # 上
                base_id + 6   # 下
            ]
        
        # 确保相邻通道ID在有效范围内
        layer = channel_id // 36  # 获取当前层号
        valid_adjacent_ids = []
        for adj_id in adjacent_ids:
            adj_layer_id = layer * 36 + adj_id
            if 0 <= adj_layer_id < len(self.subchannels):
                valid_adjacent_ids.append(adj_layer_id)
        
        return [self.subchannels[i] for i in valid_adjacent_ids] 